# -*- coding: utf-8 -*-
# 导入相应的工具包
from torch import nn
import torch
from torchvision import transforms
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
import torchvision.models as models
import matplotlib.pyplot as plt
# 导入工具包
import glob
import os
from PIL import Image
from tqdm import tqdm

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 指定数据集路径
cars_train_path = './data_cars/train/'
cars_test_path = './data_cars/test/'
# 图像增强方法
# 水平翻转
transforms_HF = transforms.RandomHorizontalFlip(p=0.5)
# 随机旋转
transforms_RR = transforms.RandomRotation(45)
# 尺寸调整
transforms_RS = transforms.Resize(size=(256, 256))
# 颜色调整
transforms_color = transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.2)
# 训练集增强方法的组合:首先转换为tensor,然后进行几何和颜色增强
train_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms_HF,
    transforms_RR,
    transforms_RS,
    transforms_color
])
# 验证集增强方法
test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms_RS,
])

# 获取训练集数据和测试集数据
cars_train = ImageFolder(cars_train_path, transform=train_transform)
cars_test = ImageFolder(cars_test_path, transform=test_transform)

# 使用dataloader获取批量数据
batch_size = 256
train_loader = DataLoader(dataset=cars_train, batch_size=batch_size, shuffle=True, pin_memory=True)
test_loader = DataLoader(dataset=cars_test, batch_size=batch_size, shuffle=False, pin_memory=True)

# 遍历每个迭代的数据，将其结果展示出来
for b, (imgs, targets) in enumerate(train_loader):
    imgs, targets = imgs.to(device), targets.to(device)
    # 获取第一个batch的图像
    if b == 1:
        # 将其进行展示
        fig, axes = plt.subplots(1, 2)
        # 遍历batch中的每个图像
        for i in range(batch_size):
            # 图像显示出来
            axes[i].imshow(imgs[i].permute(1, 2, 0))
            # 设置图像标题
            axes[i].set_title(targets[i].item())
        plt.show()
    elif b > 0:
        break

# 模型加载：进行模型微调，使用imageNet的预训练模型
model = models.resnet34(pretrained=True).to(device)
# 将所有的参数层进行冻结
for param in model.parameters():
    param.requires_grad = False
# 获取到fc层的输入
num_fc_ftr = model.fc.in_features
# 定义一个新的FC层，作为模型的输出层，进行车辆的分类
model.fc = nn.Linear(num_fc_ftr, 196).to(device)

# 学习率
learning_rate = 1e-3
# 训练轮数
num_epochs = 10
# 优化算法Adam = RMSProp + Momentum (梯度、lr两方面优化下降更快更稳)
optimizer = torch.optim.Adam(model.fc.parameters(), lr=learning_rate)
# 交叉熵损失函数
loss_fn = torch.nn.CrossEntropyLoss()


# 计算模型预测精度:测试集数据，模型
def evaluate_accuracy(data_iter, model):
    total = 0
    correct = 0
    # 不进行梯度计算
    with torch.no_grad():
        # 模型是验证模式
        model.eval()
        # 获取每一个batch的数据，进行预测
        for images, labels in data_iter:
            images = images.to(device)

            labels = labels.to(device)
            outputs = model(images)
            # 获取预测结果
            _, predicts = torch.max(outputs.data, dim=1)
            # 预测的次数
            total += labels.size(0)
            # 预测正确的个数
            correct += (predicts == labels).cpu().sum()

    # 获取准确率
    return correct / total


# 定义模型训练过程:指定数据集，优化器，损失函数和轮次
def train(data_loader=train_loader,
          optimizer=optimizer,
          loss_fn=loss_fn,
          epochs=num_epochs):
    # 遍历每一个轮次进行训练
    for epoch in range(epochs):
        print('current epoch = {}'.format(epoch))
        # 每一个轮次的损失，预测个数和预测正确个数的初始化
        train_accuracy_total = 0
        train_correct = 0
        # 损失值的和
        train_loss_sum = 0
        # 迭代次数
        iter = 0
        for i, (images, labels) in enumerate(tqdm(data_loader, desc=f"打印进度》》---epoch {epoch + 1}/{num_epochs}")):
            # 模型定义为训练模式
            images, labels = images.to(device), labels.to(device)
            model.train()
            # 对数据进行预测
            outputs = model(images)
            # 计算模型的损失
            loss = loss_fn(outputs, labels)
            # 在做反向传播前先清除网络状态
            optimizer.zero_grad()
            # 损失值进行反向传播
            loss.backward()
            # 参数迭代更新
            optimizer.step()
            # 求损失的和
            train_loss_sum += loss.item()
            # 输出模型预测结果
            _, predicts = torch.max(outputs.data, dim=1)
            # 获取训练集预测正确的个数
            train_accuracy_total += labels.size(0)
            train_correct += (predicts == labels).cpu().sum().item()
            iter += 1
        # 测试集预测的准确率
        test_acc = evaluate_accuracy(test_loader, model)
        # torch.save(model.state_dict(), "./weights/model-{}.pth".format(epoch))
        print(
            'epoch:{0},   loss:{1:.4f},   train accuracy:{2:.3f},  test accuracy:{3:.3f}'
            .format(epoch, train_loss_sum / (iter + 0.01),
                    train_correct / train_accuracy_total, test_acc))
    print('------------finish training-------------')


def predict():
    model.load_state_dict(torch.load("./model/model.pth", map_location=torch.device('cpu')))
    # 指定要预测的图像
    file_path = "./car_image/"
    # 对图像进行预测
    with torch.no_grad():
        # 指定模型为验证模式
        model.eval()
        # 获取图像的路径
        image_paths = glob.glob(os.path.join(file_path, '*.jpg'))
        # 遍历每一幅图像进行预测
        for img_path in image_paths:
            # 打开图片
            img = Image.open(img_path)
            # 增加batch维度
            batch = test_transform(img).unsqueeze(0)
            # 将图像送入网络中进行预测
            outputs = model(batch).squeeze(0).softmax(0)
            # 获取预测结果：类别索引
            class_id = outputs.argmax().item()
            # 获取预测结果：类别分数
            score = outputs[class_id].item()
            # 获取对应的类别名称
            category_name = cars_train.classes[class_id]
            # 预测结果的展示
            plt.imshow(img)
            plt.title(category_name)
            plt.show()


if __name__ == '__main__':
    train()
